import torch
from torch import nn

from einops import rearrange, repeat
from einops.layers.torch import Rearrange
import thesmuggler
ysl = thesmuggler.smuggle('/home/yusongli/Documents/vit-pytorch/_.py')

# helpers


def snoop(**kwargs):
    return pysnooper.snoop('/home/yusongli/Documents/vit-pytorch/debug.log', **kwargs)


def pair(t):
    return t if isinstance(t, tuple) else (t, t)


# classes

class FeedForward(nn.Module):
    def __init__(self, dim, hidden_dim, dropout=0.0):
        super().__init__()
        self.net = nn.Sequential(
            nn.LayerNorm(dim),
            nn.Linear(dim, hidden_dim),
            nn.GELU(),
            # nn.Dropout(dropout),
            nn.Linear(hidden_dim, dim),
            # nn.Dropout(dropout)
        )

    def forward(self, x):
        x = self.net(x)
        return x


# @snoop(watch=('x.shape', 'qkv.shape', 'q.shape', 'k.shape', 'v.shape', 'dots.shape', 'attn.shape' 'out.shape'))
class Attention(nn.Module):
    def __init__(self, dim, heads=8, dim_head=64, dropout=0.0):
        super().__init__()
        inner_dim = dim_head * heads
        project_out = heads != 1 or dim_head != dim

        self.heads = heads
        self.scale = dim_head**-0.5

        self.norm = nn.LayerNorm(dim)
        self.attend = nn.Softmax(dim = -1)
        # self.dropout = nn.Dropout(dropout)

        self.to_qkv = nn.Linear(dim, inner_dim * 3, bias=False)

        self.to_out = nn.Sequential(nn.Linear(inner_dim, dim), nn.Dropout(dropout)) if project_out else nn.Identity()

    def forward(self, x):
        x = self.norm(x)
        qkv = self.to_qkv(x).chunk(3, dim = -1)
        q, k, v = map(lambda t: rearrange(t, 'b n (h d) -> b h n d', h = self.heads), qkv)

        dots = torch.matmul(q, k.transpose(-1, -2)) * self.scale

        attn = self.attend(dots)
        # attn = self.dropout(attn)

        out = torch.matmul(attn, v)
        out = rearrange(out, 'b h n d -> b n (h d)')
        out =  self.to_out(out)
        return out

# @snoop(watch='x.shape')
class Transformer(nn.Module):
    def __init__(self, dim, depth, heads, dim_head, mlp_dim, dropout=0.0):
        super().__init__()
        self.layers = nn.ModuleList([])
        for _ in range(depth):
            self.layers.append(nn.ModuleList([
                Attention(dim, heads = heads, dim_head = dim_head, dropout = dropout),
                # FeedForward(dim, mlp_dim, dropout = dropout)
            ]))
    def forward(self, x):
        for layer in self.layers:
            for block in layer:
                import ipdb; ipdb.set_trace()  # HACK: Songli.Yu: ""
                x = block(x) + x
        return x


# @snoop(watch=('video.shape', 'x.shape', 'cls_tokens.shape'))
class ViT(nn.Module):
    def __init__(self,
                 *,
                 image_size,
                 image_patch_size,
                 frames,
                 frame_patch_size,
                 num_classes,
                 depth,
                 heads,
                 mlp_dim,
                 # pool = 'cls',
                 channels = 3,
                 # dim_head = 64,
                 dim_head = None,
                 dropout = 0.,
                 emb_dropout = 0.,
                 f=None,
                 h=None,
                 w=None,
                 dim=None):  # sourcery skip: low-code-quality
        super().__init__()
        image_height, image_width = pair(image_size)
        patch_height, patch_width = pair(image_patch_size)

        assert (
            image_height % patch_height == 0 and image_width % patch_width == 0
        ), 'Image dimensions must be divisible by the patch size.'
        assert frames % frame_patch_size == 0, 'Frames must be divisible by frame patch size'

        # num_patches = (image_height // patch_height) * (image_width // patch_width) * (frames // frame_patch_size)

        # num_patches = (f or (image_height // patch_height)) * (h or (image_width // patch_width)) * (w or (frames // frame_patch_size))
        num_patches = (h or (image_width // patch_width)) * (w or (frames // frame_patch_size))

        # patch_dim = channels * patch_height * patch_width * frame_patch_size
        patch_dim = channels * (f and (frames // f) or patch_height) * (h and (image_height // h) or patch_width) * (w and (image_width // w) or frame_patch_size)
        dim = dim or patch_dim
        dim_head = dim_head or dim // heads

        # assert pool in {'cls', 'mean'}, 'pool type must be either cls (cls token) or mean (mean pooling)'

        self.to_patch_embedding = nn.Sequential(
            # Rearrange('b c (f pf) (h p1) (w p2) -> b (f h w) (p1 p2 pf c)', p1 = patch_height, p2 = patch_width, pf = frame_patch_size),
            Rearrange(
                # 'b c (f pf) (h ph) (w pw) -> b (f h w) (ph pw pf c)',

                'b c (f pf) (h ph) (w pw) -> (b f) (h w) (ph pw pf c)',
                # 'b c (f pf) (h ph) (w pw) -> (b h) (f w) (ph pw pf c)',
                # 'b c (f pf) (h ph) (w pw) -> (b w) (f h) (ph pw pf c)',
                f=f or frames // frame_patch_size,
                h=h or image_height // patch_height,
                w=w or image_width // patch_width
            ),
            # nn.LayerNorm(patch_dim),
            # nn.Linear(patch_dim, dim),
            nn.LayerNorm(dim),
        )

        # self.pos_embedding = nn.Parameter(torch.randn(1, num_patches + 1, dim))
        self.pos_embedding = nn.Parameter(torch.randn(1, num_patches, dim))
        # self.cls_token = nn.Parameter(torch.randn(1, 1, dim))
        # self.dropout = nn.Dropout(emb_dropout)

        self.transformer = Transformer(dim, depth, heads, dim_head, mlp_dim, dropout)

        # self.pool = pool
        # self.to_latent = nn.Identity()

        # self.mlp_head = nn.Sequential(
        #     nn.LayerNorm(dim),
        #     nn.Linear(dim, num_classes)
        # )
        self.inv = nn.Sequential(
            # nn.Linear(dim, patch_dim),
            Rearrange(
                # 'b (f h w) (ph pw pf c) -> b c (f pf) (h ph) (w pw)',

                '(b f) (h w) (ph pw pf c) -> b c (f pf) (h ph) (w pw)',
                # '(b h) (f w) (ph pw pf c) -> b c (f pf) (h ph) (w pw)',
                # '(b w) (f h) (ph pw pf c) -> b c (f pf) (h ph) (w pw)',

                f=f or (frames // frame_patch_size),
                h=h or (image_height // patch_height),
                w=w or (image_width // patch_width),
                ph=h and (image_height // h) or patch_height,
                pw=w and (image_width // w) or patch_width,
                pf=f and (frames // f) or frame_patch_size
            )
        )

    # def forward(self, video):
    def forward(self, x):
        # x = self.to_patch_embedding(video)
        x = self.to_patch_embedding(x)
        b, n, _ = x.shape

        # cls_tokens = repeat(self.cls_token, '1 1 d -> b 1 d', b = b)
        # x = torch.cat((cls_tokens, x), dim=1)
        # x += self.pos_embedding[:, :(n + 1)]
        x += self.pos_embedding
        # x = self.dropout(x)

        x = self.transformer(x)

        # x = x.mean(dim = 1) if self.pool == 'mean' else x[:, 0]

        # x = self.to_latent(x)
        # x = self.mlp_head(x)

        x = self.inv(x)
        return x


if __name__ == '__main__':
    # shapes1 = (2, 32, 32, 80, 96)
    # shapes = (2, 32, 32, 80, 96)

    shapes1 = (2, 32, 32, 80, 96)
    shapes = (2, 64, 32, 40, 48)

    # shapes1 = (2, 64, 32, 40, 48)
    # shapes = (2, 128, 32, 20, 24)

    # shapes1 = (2, 128, 32, 20, 24)
    # shapes = (2, 256, 16, 10, 12)

    patches = (4, 10, 12)

    v = ViT(
        image_size=shapes[-2:],
        image_patch_size=patches[-2:],
        frames=shapes[-3],
        frame_patch_size=patches[-3],
        num_classes=1000,
        # dim=1024,
        depth=1,
        heads=8,
        mlp_dim=2048,
        #
        # pool = 'cls',
        channels = shapes[1],
        # dim_head = 64,
        dropout = 0.,
        emb_dropout = 0.,
        #
        f=shapes1[-3] // patches[-3],
        h=shapes1[-2] // patches[-2],
        w=shapes1[-1] // patches[-1],
    )
    video = torch.ones(*shapes) # (batch, channels, frames, height, width)

    preds = ysl.all(v)(video)
    print(preds.shape)
